// Package middleware provides an in-memory token bucket rate limiter.
package middleware

import (
	"net"
	"net/http"
	"sync"
	"time"
)

// In-memory token bucket per client (ip) and route prefix
type bucket struct {
	tokens float64
	last   time.Time
}

type RateLimiter struct {
	mu      sync.Mutex
	buckets map[string]*bucket
	rate    float64 // tokens per second
	burst   float64
}

func NewRateLimiter(rate float64, burst int) *RateLimiter {
	return &RateLimiter{buckets: map[string]*bucket{}, rate: rate, burst: float64(burst)}
}

func (rl *RateLimiter) key(r *http.Request) string {
	host, _, _ := net.SplitHostPort(r.RemoteAddr)
	if host == "" {
		host = r.RemoteAddr
	}
	return host + "|" + r.URL.Path
}

func (rl *RateLimiter) Allow(r *http.Request) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	k := rl.key(r)
	b := rl.buckets[k]
	now := time.Now()
	if b == nil {
		b = &bucket{tokens: rl.burst, last: now}
		rl.buckets[k] = b
	}
	// replenish
	elapsed := now.Sub(b.last).Seconds()
	b.tokens += elapsed * rl.rate
	if b.tokens > rl.burst {
		b.tokens = rl.burst
	}
	b.last = now
	if b.tokens >= 1 {
		b.tokens -= 1
		return true
	}
	return false
}

func (rl *RateLimiter) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !rl.Allow(r) {
			w.WriteHeader(http.StatusTooManyRequests)
			_, _ = w.Write([]byte("rate limit exceeded"))
			return
		}
		next.ServeHTTP(w, r)
	})
}
